home *** CD-ROM | disk | FTP | other *** search
/ Enter 2006 September / Enter 09 2006.iso / Internet / SpamExperts Home 1.1 / SpamExperts Home.exe / lib / spamexperts.modules / spamexperts / imap4.pyc (.txt) < prev    next >
Encoding:
Python Compiled Bytecode  |  2006-07-14  |  35.2 KB  |  1,117 lines
  1. # Source Generated with Decompyle++
  2. # File: in.pyc (Python 2.4)
  3.  
  4. import os
  5. import re
  6. import sys
  7. import md5
  8. import time
  9. import copy
  10. import email
  11. import base64
  12. import socket
  13. import string
  14. import thread
  15. import imaplib
  16. import traceback
  17. import threading
  18. import email.Header as email
  19.  
  20. try:
  21.     frozenset
  22. except NameError:
  23.     from sets import ImmutableSet as frozenset
  24.  
  25. from spambayes import Dibbler
  26. from spambayes import storage
  27. from spambayes.message import insert_exception_header, PERSISTENT_HAM_STRING
  28. from spamexperts import Options
  29. verbose = Options.options[('globals', 'verbose')]
  30. Options.options[('globals', 'verbose')] = False
  31. from spambayes.scripts.sb_server import ServerLineReader
  32. Options.options[('globals', 'verbose')] = verbose
  33. del verbose
  34. from spamexperts import ProxyClassifier
  35. from spamexperts.message import SEHeaderMessage
  36. from se_config import spamexpertsConfig as configuration
  37. from spamexperts.OptionsClass import IS_HAM, IS_SPAM, BLOCKED, IS_UNSURE
  38.  
  39. class IMAP4ProxyBase(Dibbler.BrighterAsyncChat):
  40.     """An async dispatcher that understands IMAP4 and proxies to an IMAP4
  41.     server, calling `self.onTransaction(request, response)` for each
  42.     transaction.  Generally similar to the POP3ProxyBase class.
  43.  
  44.     self.onTransaction() should return the response to pass back to the
  45.     email client - the response can be the verbatim response or a processed
  46.     version of it.  The special command 'KILL' kills it (passing a 'LOGOUT'
  47.     command to the server).
  48.     """
  49.     
  50.     def __init__(self, clientSocket, serverName, serverPort, ssl = False):
  51.         Dibbler.BrighterAsyncChat.__init__(self, clientSocket)
  52.         self.request = ''
  53.         self.response = ''
  54.         self.set_terminator('\r\n')
  55.         self.command = ''
  56.         self.command_id = None
  57.         self.args = []
  58.         self.isClosing = False
  59.         if not self.onIncomingConnection(clientSocket):
  60.             self.push('-ERR Connection not allowed\r\n')
  61.             self.close_when_done()
  62.             return None
  63.         
  64.         self.serverName = serverName
  65.         self.serverPort = serverPort
  66.         self.serverSocket = ServerLineReader(serverName, serverPort, self.onServerLine, ssl)
  67.  
  68.     
  69.     def close(self):
  70.         if hasattr(self, 'serverSocket'):
  71.             self.serverSocket.close()
  72.         
  73.         Dibbler.BrighterAsyncChat.close(self)
  74.  
  75.     
  76.     def onIncomingConnection(self, clientSocket):
  77.         '''Checks the security settings.'''
  78.         remoteIP = clientSocket.getpeername()[0]
  79.         trustedIPs = Options.options[('imap4proxy', 'allow_remote_connections')]
  80.         if trustedIPs == '*' or remoteIP == clientSocket.getsockname()[0]:
  81.             return True
  82.         
  83.         trustedIPs = trustedIPs.replace('.', '\\.').replace('*', '([01]?\\d\\d?|2[04]\\d|25[0-5])')
  84.         for trusted in trustedIPs.split(','):
  85.             if re.search('^' + trusted + '$', remoteIP):
  86.                 return True
  87.                 continue
  88.         
  89.         return False
  90.  
  91.     
  92.     def onTransaction(self, command_id, command, args, response):
  93.         '''Overide this.  Takes the raw request and the response, and
  94.         returns the (possibly processed) response to pass back to the
  95.         email client.
  96.         '''
  97.         raise NotImplementedError
  98.  
  99.     
  100.     def onServerLine(self, line):
  101.         '''A line of response has been received from the IMAP4 server.'''
  102.         self.response = self.response + line
  103.         if not line:
  104.             self.isClosing = True
  105.             self.onResponse()
  106.             self.response = ''
  107.             return None
  108.         
  109.         if not self.command:
  110.             self.push(self.response)
  111.             self.response = ''
  112.         elif line.startswith('+ '):
  113.             self.push(self.response)
  114.             self.response = ''
  115.         elif line.startswith(self.command_id):
  116.             self.onResponse()
  117.             self.response = ''
  118.         
  119.  
  120.     
  121.     def collect_incoming_data(self, data):
  122.         '''Asynchat override.'''
  123.         self.request = self.request + data
  124.  
  125.     
  126.     def found_terminator(self):
  127.         '''Asynchat override.'''
  128.         self.serverSocket.push(self.request + '\r\n')
  129.         if self.request.strip() == '':
  130.             self.command = ''
  131.             self.args = []
  132.         else:
  133.             splitCommand = self.request.strip().split()
  134.             if len(splitCommand) < 3:
  135.                 self.args += tuple(splitCommand)
  136.             else:
  137.                 self.command_id = splitCommand[0]
  138.                 self.command = splitCommand[1].upper()
  139.                 self.args = splitCommand[2:]
  140.                 if self.command == 'UID':
  141.                     self.uid = True
  142.                     self.command = self.args[0].upper()
  143.                     self.args = self.args[1:]
  144.                 else:
  145.                     self.uid = False
  146.         self.request = ''
  147.  
  148.     
  149.     def onResponse(self):
  150.         if self.response:
  151.             cooked = self.onTransaction(self.command_id, self.command, self.args, self.response)
  152.             self.push(cooked)
  153.         
  154.         if self.isClosing:
  155.             self.close_when_done()
  156.         
  157.         self.command_id = None
  158.         self.command = ''
  159.         self.args = []
  160.         self.isClosing = False
  161.  
  162.  
  163.  
  164. class SEIMAP4Proxy(IMAP4ProxyBase, ProxyClassifier.ProxyClassifier):
  165.     '''Proxies between an email client and an IMAP4 server, inserting
  166.     judgement headers.  It acts on the following IMAP4 commands:
  167.      o FETCH:
  168.         o If the component(s) being fetched is(are) message(s) then
  169.           adds the judgement header based on the raw headers and body
  170.           of the message.
  171.      o LOGIN:
  172.         o Does no processing based on the LOGIN command itself, but
  173.           expires any old messages in the two caches.
  174.      o AUTHENTICATE:
  175.         o Same as LOGIN.
  176.     '''
  177.     
  178.     def __init__(self, clientSocket, serverName, serverPort, state, ssl = False):
  179.         self.skip_logging = None
  180.         self.authenticating = None
  181.         self.isClosed = False
  182.         ProxyClassifier.ProxyClassifier.__init__(self)
  183.         IMAP4ProxyBase.__init__(self, clientSocket, serverName, serverPort, ssl)
  184.         self.handlers = {
  185.             'FETCH': self.onFetch,
  186.             'LOGIN': self.onLogin,
  187.             'AUTHENTICATE': self.onLogin }
  188.         self.state = state
  189.         self.state.totalSessions += 1
  190.         self.state.activeSessions += 1
  191.         self.add_when_cooking = True
  192.         self.state.proxies.append(self)
  193.  
  194.     
  195.     def send(self, data):
  196.         '''Logs the data to the log file.'''
  197.         if self.skip_logging:
  198.             self.state.imapLogFile.write('...FETCH CONTENTS...\r\n')
  199.         else:
  200.             self.state.imapLogFile.write(data)
  201.         if self.skip_logging:
  202.             split_data = data.split('\n')
  203.             if len(split_data) == 1 and split_data[0].startswith(self.skip_logging):
  204.                 self.skip_logging = False
  205.             elif len(split_data) > 1 and split_data[-2].startswith(self.skip_logging):
  206.                 self.skip_logging = False
  207.             
  208.         
  209.         self.state.logFile.flush()
  210.         
  211.         try:
  212.             return IMAP4ProxyBase.send(self, data)
  213.         except socket.error:
  214.             self.close()
  215.  
  216.         return 0
  217.  
  218.     
  219.     def recv(self, size):
  220.         '''Logs the data to the log file.'''
  221.         
  222.         try:
  223.             data = IMAP4ProxyBase.recv(self, size)
  224.         except socket.error:
  225.             e = None
  226.             if e.args[0] == 10053:
  227.                 pass
  228.             else:
  229.                 print >>sys.stderr, 'Unexpected socket error:', str(e)
  230.                 if Options.options[('globals', 'verbose')]:
  231.                     traceback.print_exc(None, sys.stderr)
  232.                 
  233.             self.close()
  234.             return ''
  235.  
  236.         check_data = data.lower().split()
  237.         if hasattr(self, 'authenticating') and self.authenticating and check_data:
  238.             if check_data[0] == self.authenticating:
  239.                 self.authenticating = None
  240.                 log_data = data
  241.             else:
  242.                 log_data = 'XXX AUTHENTICATION DETAILS XXX\r\n'
  243.         elif len(check_data) > 1 and check_data[1].startswith('login'):
  244.             log_data = check_data[0] + ' LOGIN XXXXXXXX XXXXXXXX\r\n'
  245.         elif len(check_data) > 1 and check_data[1].startswith('authenticate'):
  246.             self.authenticating = check_data[0]
  247.             log_data = data
  248.         elif (len(check_data) > 1 or check_data[1].startswith('fetch') or len(check_data) > 2) and check_data[1].startswith('uid') and check_data[2].startswith('fetch'):
  249.             self.skip_logging = check_data[0]
  250.             log_data = data
  251.         else:
  252.             log_data = data
  253.         self.state.imapLogFile.write(log_data)
  254.         self.state.imapLogFile.flush()
  255.         return data
  256.  
  257.     
  258.     def close(self):
  259.         if not self.isClosed:
  260.             self.isClosed = True
  261.             self.state.activeSessions -= 1
  262.             IMAP4ProxyBase.close(self)
  263.             self.serverSocket.close()
  264.             self.state.proxies.remove(self)
  265.         
  266.  
  267.     
  268.     def onTransaction(self, command_id, command, args, response):
  269.         '''Takes the raw request and response, and returns the
  270.         (possibly processed) response to pass back to the email client.
  271.         '''
  272.         handler = self.handlers.get(command, self.onUnknown)
  273.         return handler(command_id, command, args, response)
  274.  
  275.     find_header = re.compile('(\\d+)\\s+fetch.*(?:BODY(?:\\.PEEK)?\\[HEADER\\])\\s+{(\\d+)\\}\\r?\\n(.*)', re.DOTALL | re.IGNORECASE)
  276.     find_body = re.compile('fetch.*(?:BODY(?:\\.PEEK)?\\[\\]|RFC822)\\s+\\{(\\d+)\\}\\r?\\n(.*)', re.DOTALL | re.IGNORECASE)
  277.     find_size = re.compile('fetch.*(rfc822.size\\s+)(\\d*)', re.DOTALL | re.IGNORECASE)
  278.     replace_size = re.compile('rfc822.size\\s+(\\d*)', re.DOTALL | re.IGNORECASE)
  279.     
  280.     def onFetch(self, unused, unused2, unused3, response):
  281.         '''If any messages are being fetched, then adds the judgement
  282.         header based on the raw headers and body of the message(s).  If
  283.         something else (e.g. flags) are being fetched, then just proxies
  284.         straight through.'''
  285.         self.state.model_notifier.SetBeginUpdating()
  286.         mo = self.find_body.search(response)
  287.         if mo:
  288.             literal_size = int(mo.group(1))
  289.             literal = mo.group(2)[:literal_size]
  290.             new_version = self.body_cooker(literal)
  291.             response = response.replace(literal, new_version)
  292.             response = response.replace('{%d}' % (literal_size,), '{%d}' % (len(new_version),))
  293.         
  294.         mo = self.find_size.search(response)
  295.         if mo:
  296.             rfc822_wording = mo.group(1)
  297.             rfc822_size = mo.group(2)
  298.             new_version = '%s%s' % (rfc822_wording, self.size_cooker(rfc822_size))
  299.             response = self.replace_size.sub(new_version, response)
  300.         
  301.         if configuration.eudora_compatibility:
  302.             mo = self.find_header.search(response)
  303.             if mo:
  304.                 msg_id = mo.group(1)
  305.                 literal_size = int(mo.group(2))
  306.                 literal = mo.group(3)[:literal_size]
  307.                 new_version = self.header_cooker(msg_id, literal)
  308.                 response = response.replace(literal, new_version)
  309.                 response = response.replace('{%d}' % (literal_size,), '{%d}' % (len(new_version),))
  310.             
  311.         
  312.         self.state.model_notifier.SetEndUpdating()
  313.         return response
  314.  
  315.     
  316.     def size_cooker(self, old_size):
  317.         return str(int(old_size) + self.HEADER_SIZE_FUDGE_FACTOR)
  318.  
  319.     
  320.     def generate_id(self, messageText):
  321.         return str(int(md5.md5(messageText).hexdigest(), 16))
  322.  
  323.     split_header_re = re.compile('\\n\\r?\\n')
  324.     
  325.     def header_cooker(self, msg_id, headerText):
  326.         
  327.         try:
  328.             self.serverSocket.send('sehc FETCH %s (RFC822)\n' % (msg_id,))
  329.         except socket.error:
  330.             e = None
  331.             print >>sys.stderr, 'Unexpected socket error:', str(e)
  332.             if Options.options[('globals', 'verbose')]:
  333.                 traceback.print_exc(None, sys.stderr)
  334.             
  335.             return headerText
  336.  
  337.         response = []
  338.         while True:
  339.             
  340.             try:
  341.                 data = self.serverSocket.recv(1024)
  342.             except socket.error:
  343.                 e = None
  344.                 if e[0] == 10035:
  345.                     continue
  346.                 
  347.                 print >>sys.stderr, 'Unexpected socket error:', str(e)
  348.                 if Options.options[('globals', 'verbose')]:
  349.                     traceback.print_exc(None, sys.stderr)
  350.                 
  351.                 return headerText
  352.  
  353.             response.append(data)
  354.             if '\nsehc' in data:
  355.                 break
  356.                 continue
  357.         response = ''.join(response)
  358.         temp = response[response.find('RFC822 {') + 8:]
  359.         temp2 = temp.find('}')
  360.         length = int(temp[:temp2])
  361.         response = temp[temp2 + 3:temp2 + 2 + length]
  362.         cooked = self.body_cooker(response)
  363.         return self.split_header_re.split(cooked, 1)[0] + '\n\r\n'
  364.  
  365.     
  366.     def body_cooker(self, messageText):
  367.         
  368.         try:
  369.             msg = email.message_from_string(messageText, _class = SEHeaderMessage)
  370.             msg.setId(self.generate_id(messageText))
  371.             (score, classification, clues) = self.classify_message(msg)
  372.             if classification == IS_HAM:
  373.                 classification = Options.options[('Headers', 'header_ham_string')]
  374.             elif classification == IS_SPAM:
  375.                 classification = Options.options[('Headers', 'header_spam_string')]
  376.             elif classification == IS_UNSURE:
  377.                 classification = Options.options[('Headers', 'header_unsure_string')]
  378.             
  379.             isTooBig = False
  380.             isSuppressedBulkHam = False
  381.             if not configuration.block_spam:
  382.                 if classification == Options.options[('Headers', 'header_ham_string')] and Options.options[('Storage', 'no_cache_bulk_ham')]:
  383.                     pass
  384.                 isSuppressedBulkHam = msg.get('precedence') in [
  385.                     'bulk',
  386.                     'list']
  387.                 size_limit = Options.options[('Storage', 'no_cache_large_messages')]
  388.                 if size_limit > 0:
  389.                     pass
  390.                 isTooBig = len(messageText) > size_limit
  391.             
  392.             msg.RememberClassification(classification)
  393.             msg.addHeaders(prob = score, clues = clues)
  394.             tid = msg.getId()
  395.             if self.add_when_cooking and Options.options[('Storage', 'cache_messages')] and not isSuppressedBulkHam and not isTooBig:
  396.                 if classification == Options.options[('Headers', 'header_spam_string')]:
  397.                     self.state.numSpams += 1
  398.                     print >>sys.stderr, 'Store msg %s in Spamcorpus %s' % (tid, classification)
  399.                     corpus = self.state.spamCorpus
  400.                 elif classification == Options.options[('Headers', 'header_unsure_string')]:
  401.                     print >>sys.stderr, 'Store msg %s in Unsurecorpus %s' % (tid, classification)
  402.                     corpus = self.state.unsureCorpus
  403.                 else:
  404.                     self.state.numHams += 1
  405.                     print >>sys.stderr, 'Store msg %s in Hamcorpus %s' % (tid, classification)
  406.                     corpus = self.state.hamCorpus
  407.                 message = corpus.makeMessage(msg.getId(), msg.as_string())
  408.                 corpus.addMessage(message)
  409.             
  410.             messageText = msg.as_string()
  411.         except:
  412.             (messageText, details) = insert_exception_header(messageText)
  413.             print >>sys.stderr, details
  414.  
  415.         return messageText
  416.  
  417.     
  418.     def onLogin(self, unused, unused2, args, response, do_welcome = True):
  419.         '''Spins off two separate threads that expires any old messages
  420.         in the caches, but does not do any processing of the LOGIN
  421.         command itself.'''
  422.         thread.start_new_thread(self.state.spamCorpus.removeExpiredMessages, ())
  423.         thread.start_new_thread(self.state.hamCorpus.removeExpiredMessages, ())
  424.         if response.split()[1].lower() == 'ok':
  425.             self.current_account = '%s_%s_IMAP' % (args[0], self.serverName)
  426.             if self.current_account not in self.state.delayed_messages and do_welcome:
  427.                 self.state.delayed_messages[self.current_account] = { }
  428.                 self.state.delayed_messages.store()
  429.                 welcomeText = self.get_welcome_message(args[0], self.serverName)
  430.                 self.serverSocket.setblocking(True)
  431.                 self.serverSocket.send('welcome APPEND INBOX {%d}\r\n' % (len(welcomeText),))
  432.                 self.serverSocket.recv(1024)
  433.                 self.serverSocket.send(welcomeText + '\r\n')
  434.                 self.serverSocket.recv(1024)
  435.                 self.serverSocket.setblocking(False)
  436.             
  437.         
  438.         return response
  439.  
  440.     
  441.     def onUnknown(self, unused, unused2, unused3, response):
  442.         """Default handler; returns the server's response verbatim."""
  443.         return response
  444.  
  445.  
  446.  
  447. class IMAPFilter(object):
  448.     initial_key = 'initial_access'
  449.     
  450.     def close(self):
  451.         
  452.         try:
  453.             self.state.proxies.remove(self)
  454.         except ValueError:
  455.             pass
  456.  
  457.  
  458.     
  459.     def load_folders_to_filter(self):
  460.         folders = _[1]
  461.         self.filter_folders = tuple(folders)
  462.  
  463.     terminated = False
  464.     
  465.     def updateMessageDatabase(self, imap_connection):
  466.         """Run through the IMAP server (the known folders) and update our
  467.         local information about whether messages are spam or not.  This
  468.         will take a considerable amount of time, so we cache this
  469.         information, and don't reconsider messages we have already looked
  470.         at.
  471.  
  472.         Download a copy (and delete from the server) any spam messages,
  473.         and send these to the 'blocked' database.
  474.  
  475.         Perhaps we should reclassify all messages whenever the database
  476.         changes - after all, it's possible that the classification will
  477.         change.  This might mean releasing messages automatically (since
  478.         we took them off the server) and would be complicated and
  479.         time-consuming.  For now, we do not do this.
  480.         """
  481.         i = imap_connection
  482.         for folder in self.filter_folders:
  483.             
  484.             try:
  485.                 self.updateMessageDatabaseFolder(i, folder)
  486.             continue
  487.             except Exception:
  488.                 if self.terminated:
  489.                     break
  490.                 
  491.                 raise 
  492.                 continue
  493.             
  494.  
  495.         
  496.         i.logout()
  497.         
  498.         try:
  499.             self.state.open_remote_connections.remove(self.current_account)
  500.         except ValueError:
  501.             None<EXCEPTION MATCH>Exception
  502.             None<EXCEPTION MATCH>Exception
  503.             print >>sys.stderr, 'Connection was missing:', self.current_account
  504.         except:
  505.             None<EXCEPTION MATCH>Exception
  506.  
  507.  
  508.     
  509.     class _dummy_msg(object):
  510.         '''Dummy message class used to check if a message is in the
  511.         message info database.'''
  512.         
  513.         def __init__(self, key):
  514.             
  515.             self.getDBKey = lambda : key
  516.             self.stored_attributes = [
  517.                 'c',
  518.                 't',
  519.                 'block_state',
  520.                 'account',
  521.                 'date_modified',
  522.                 'internaldate',
  523.                 'flags',
  524.                 'folder_name',
  525.                 'uid']
  526.             for att in self.stored_attributes:
  527.                 setattr(self, att, None)
  528.             
  529.  
  530.  
  531.     uid_re = re.compile('(?P<id>\\d+) \\(UID (?P<data>\\d+)\\)', re.IGNORECASE)
  532.     flags_re = re.compile('(?P<id>\\d+) \\(FLAGS \\((?P<data>.*?)\\)\\)', re.IGNORECASE)
  533.     internaldate_re = re.compile('(?P<id>\\d+) \\(INTERNALDATE \\"(?P<data>.+)\\"\\)', re.IGNORECASE)
  534.     
  535.     def _parseFetchResponse(self, msg_id, regex, response):
  536.         mo = regex.search(response)
  537.         if mo:
  538.             if msg_id != mo.group('id'):
  539.                 print >>sys.stderr, 'Error in response (%s != %s)' % (msg_id, mo.group('id'))
  540.                 return None
  541.             
  542.             return mo.group('data')
  543.         
  544.         print >>sys.stderr, 'Error in response: %s, %s' % (msg_id, response)
  545.  
  546.     
  547.     def _fetchAttribute(self, imap_connection, msg_id, attribute, parser, log = True, uid = False):
  548.         
  549.         try:
  550.             if uid:
  551.                 cmd = imap_connection.uid
  552.                 args = ('FETCH', msg_id, '(%s)' % (attribute,))
  553.             else:
  554.                 cmd = imap_connection.fetch
  555.                 args = (msg_id, '(%s)' % (attribute,))
  556.             (unused, att) = cmd(*args)
  557.         except (imaplib.IMAP4.error, socket.error):
  558.             e = None
  559.             print >>sys.stderr, 'Cannot get', attribute, msg_id
  560.             print >>sys.stderr, str(e)
  561.             return None
  562.  
  563.         if log:
  564.             self.state.imapLogFile.write('FETCH %s %s: %s\r\n' % (msg_id, attribute, att))
  565.         else:
  566.             self.state.imapLogFile.write('...DATA...\r\n')
  567.         self.state.imapLogFile.flush()
  568.         if parser is not None:
  569.             return self._parseFetchResponse(msg_id, parser, att[0])
  570.         
  571.         return att
  572.  
  573.     
  574.     def _normalise_name(self, name):
  575.         return [](_[1])
  576.  
  577.     
  578.     def _get_messageinfo_database_key(self, folder_name, uid):
  579.         norm_current_account = self._normalise_name(self.current_account)
  580.         norm_folder_name = self._normalise_name(folder_name)
  581.         norm_uid = self._normalise_name(uid)
  582.         return '%s_%s_%s' % (norm_current_account, norm_folder_name.lower(), norm_uid.lower())
  583.  
  584.     
  585.     def updateMessageDatabaseFolder(self, imap_connection, folder_name):
  586.         '''Run through folder and update message status in database.'''
  587.         i = imap_connection
  588.         
  589.         try:
  590.             (unused, unused2) = i.select(folder_name, True)
  591.         except imaplib.IMAP4.error:
  592.             e = None
  593.             print >>sys.stderr, 'Cannot select folder', folder_name
  594.             print >>sys.stderr, str(e)
  595.             return None
  596.  
  597.         self.state.imapLogFile.write('Examine %s\r\n' % (folder_name,))
  598.         self.state.imapLogFile.flush()
  599.         
  600.         try:
  601.             (unused, ids) = i.search(None, 'UNDELETED')
  602.         except imaplib.IMAP4.error:
  603.             e = None
  604.             print >>sys.stderr, 'Cannot search for undeleted messages in folder', folder_name
  605.             print >>sys.stderr, str(e)
  606.             return None
  607.  
  608.         self.state.imapLogFile.write('Search UNDELETED: %s\r\n' % (ids,))
  609.         self.state.imapLogFile.flush()
  610.         for msg_id in ids[0].split():
  611.             if self.terminated:
  612.                 return None
  613.             
  614.             uid = self._fetchAttribute(i, msg_id, 'UID', self.uid_re)
  615.             if uid is None:
  616.                 continue
  617.             
  618.             flags = self._fetchAttribute(i, msg_id, 'FLAGS', self.flags_re)
  619.             if flags is None:
  620.                 continue
  621.             
  622.             db_key = self._get_messageinfo_database_key(folder_name, uid)
  623.             msg = self._dummy_msg(db_key)
  624.             self.state.message_info_database.load_msg(msg)
  625.             if msg.c:
  626.                 internaldate = [
  627.                     None]
  628.                 messageText = [
  629.                     None]
  630.             else:
  631.                 internaldate = self._fetchAttribute(i, msg_id, 'INTERNALDATE', self.internaldate_re)
  632.                 if internaldate is None:
  633.                     continue
  634.                 
  635.                 time_tuple = imaplib.Internaldate2tuple('INTERNALDATE "' + internaldate + '"')
  636.                 if time_tuple:
  637.                     msg_time = time.mktime(time_tuple)
  638.                 else:
  639.                     msg_time = None
  640.                 
  641.                 try:
  642.                     initial_key = self.state.blocked_messages[self.current_account][self.initial_key]
  643.                 except KeyError:
  644.                     e = None
  645.                     if Options.options[('globals', 'verbose')]:
  646.                         print "Skipping, because don't know", e
  647.                         continue
  648.                     continue
  649.  
  650.                 if msg_time < initial_key:
  651.                     continue
  652.                 
  653.                 messageText = self._fetchAttribute(i, msg_id, 'BODY.PEEK[]', None, log = False)
  654.                 if messageText is None:
  655.                     continue
  656.                 
  657.             self.updateMessageDatabaseMessage(uid, flags, internaldate, messageText[0], folder_name)
  658.         
  659.  
  660.     
  661.     def updateMessageDatabaseMessage(self, uid, flags, internaldate, message_text, folder_name):
  662.         '''Update message database for this message.'''
  663.         db_key = self._get_messageinfo_database_key(folder_name, uid)
  664.         msg = self._dummy_msg(db_key)
  665.         self.state.message_info_database.load_msg(msg)
  666.         old = [](_[1])
  667.         msg.uid = uid
  668.         msg.flags = flags
  669.         msg.folder_name = folder_name
  670.         msg.internaldate = internaldate
  671.         for att in msg.stored_attributes:
  672.             if att not in old or getattr(msg, att) != old[att]:
  673.                 self.state.message_info_database.store_msg(msg)
  674.                 break
  675.                 continue
  676.             []
  677.         
  678.         if message_text:
  679.             message_text = message_text[1]
  680.         else:
  681.             return None
  682.         self.processing_queue.put((message_text, db_key, { }, self.current_account))
  683.  
  684.     
  685.     def expungeMessages(self, connection):
  686.         '''Delete mail that the user has asked to be deleted.'''
  687.         leftover = { }
  688.         messages_to_delete = self.state.delete_messages[self.current_account]
  689.         to_delete = copy.copy(messages_to_delete.keys())
  690.         for msg_id in to_delete:
  691.             msg = self._dummy_msg(msg_id)
  692.             self.state.message_info_database.load_msg(msg)
  693.             if msg.uid == -1:
  694.                 if Options.options[('globals', 'verbose')]:
  695.                     print 'Skipping', msg_id, "(we don't know the UID)"
  696.                     continue
  697.                 continue
  698.             
  699.             if Options.options[('globals', 'verbose')]:
  700.                 print 'Deleting message from server (id %s uid %s)' % (msg_id, msg.uid)
  701.             
  702.             connection.select(msg.folder_name)
  703.             
  704.             try:
  705.                 connection.uid('STORE', msg.uid, '+FLAGS.SILENT', '(\\Deleted \\Seen)')
  706.             continue
  707.             except imaplib.IMAP4.error:
  708.                 e = None
  709.                 print >>sys.stderr, 'Error occured while deleting.'
  710.                 print >>sys.stderr, str(e)
  711.                 leftover[msg_id] = { }
  712.                 continue
  713.                 continue
  714.             
  715.  
  716.         
  717.         self.state.delete_messages[self.current_account] = leftover
  718.         self.state.delete_messages.store()
  719.  
  720.  
  721. CHUNK_SIZE = 2048
  722.  
  723. class ChunkingIMAP4(imaplib.IMAP4):
  724.     '''Try to avoid MemoryErrors by reading from the
  725.     socket in smaller chunks, and then combining.'''
  726.     
  727.     def read(self, size):
  728.         """Read 'size' bytes from remote."""
  729.         buffer = []
  730.         for unused in xrange(size // CHUNK_SIZE):
  731.             buffer.append(self.file.read(CHUNK_SIZE))
  732.         
  733.         buffer.append(self.file.read(size % CHUNK_SIZE))
  734.         buffer = ''.join(buffer)
  735.         if not len(buffer) == size:
  736.             raise AssertionError, 'Chunking is wrong'
  737.         return buffer
  738.  
  739.  
  740.  
  741. class ChunkingIMAP4_SSL(imaplib.IMAP4_SSL):
  742.     '''Try to avoid MemoryErrors by reading from the
  743.     socket in smaller chunks, and then combining.'''
  744.     
  745.     def read(self, size):
  746.         """Read 'size' bytes from remote."""
  747.         chunks = []
  748.         read = 0
  749.         while read < size:
  750.             data = self.sslobj.read(size - read % CHUNK_SIZE)
  751.             read += len(data)
  752.             chunks.append(data)
  753.         return ''.join(chunks)
  754.  
  755.  
  756.  
  757. class SEBlockingIMAP4Proxy(SEIMAP4Proxy, IMAPFilter):
  758.     '''Like the parent class, except that instead of marking messages
  759.     and letting them through, only ham is let through and spam is
  760.     blocked for later release or deletion.
  761.  
  762.     Essentially:
  763.  
  764.      o On connection, spin off a thread to filter mailboxes (much as
  765.        the SpamBayes IMAP filter works), recording the spam status of each
  766.        message (skipping any that have already been filtered, obviously).
  767.        Also download a copy of (and delete from the server) any spam
  768.        messages to the "blocked" database.
  769.  
  770.      o Modify FETCH responses to include the classification headers.
  771.        This behaviour isn\'t necessary for the software to work, so is not
  772.        currently implemented.
  773.  
  774.      o STORE any user-identified (via the GUI) false positives back on the
  775.        server.
  776.  
  777.      o Monitor EXAMINE and SELECT commands to keep track of which mailboxes
  778.        are used.  This way we can ignore any non-mail mailboxes, such as
  779.        are common with poorly configured IMAP servers.
  780.  
  781.      o Proxy any other commands straight through.
  782.     '''
  783.     
  784.     def __init__(self, clientSocket, serverName, serverPort, state, ssl = False):
  785.         self.skip_logging = None
  786.         self.isClosed = False
  787.         ProxyClassifier.ProxyClassifier.__init__(self)
  788.         IMAP4ProxyBase.__init__(self, clientSocket, serverName, serverPort, ssl)
  789.         IMAPFilter.__init__(self)
  790.         self.handlers = {
  791.             'LOGIN': self.onLogin,
  792.             'AUTHENTICATE': self.onAuthenticate,
  793.             'EXAMINE': self.onSelect,
  794.             'SELECT': self.onSelect,
  795.             'FETCH': self.onFetch }
  796.         self.add_when_cooking = False
  797.         self.threads = []
  798.         self.state = state
  799.         self.ssl = ssl
  800.         self.state.totalSessions += 1
  801.         self.state.activeSessions += 1
  802.         self.state.proxies.append(self)
  803.  
  804.     
  805.     def close(self):
  806.         if not self.isClosed:
  807.             self.isClosed = True
  808.             IMAP4ProxyBase.close(self)
  809.             thread.start_new_thread(self._wait_for_threads, ())
  810.         
  811.  
  812.     
  813.     def _wait_for_threads(self):
  814.         '''Let the state know that we are completely done once all the
  815.         threads are done.'''
  816.         for update_thread in self.threads:
  817.             if update_thread.isAlive():
  818.                 update_thread.join()
  819.                 continue
  820.         
  821.         self.threads = []
  822.         self.state.activeSessions -= 1
  823.         self.state.proxies.remove(self)
  824.  
  825.     
  826.     def _setupAccount(self, username):
  827.         self.current_account = '%s_%s_IMAP' % (username, self.serverName)
  828.         is_new = False
  829.         if self.current_account not in self.state.blocked_messages:
  830.             ca = self.current_account
  831.             self.state.blocked_messages[ca] = { }
  832.             ik = self.initial_key
  833.             self.state.blocked_messages[ca][ik] = time.time() + 120
  834.             self.state.blocked_messages.store()
  835.             is_new = True
  836.         
  837.         if self.current_account not in self.state.delete_messages:
  838.             self.state.delete_messages[self.current_account] = { }
  839.             self.state.delete_messages.store()
  840.             is_new = True
  841.         
  842.         if self.current_account not in self.state.delayed_messages:
  843.             self.state.delayed_messages[self.current_account] = { }
  844.             self.state.delayed_messages.store()
  845.             is_new = True
  846.         
  847.         return is_new
  848.  
  849.     
  850.     def onAuthenticate(self, command_id, command, args, response):
  851.         """Spins off two separate threads that expires any old messages
  852.         in the caches, and another thread to update the message database.
  853.         Also creates an account (to hold the blocked messages) if the
  854.         server has not been seen before, and release any 'delayed'
  855.         (i.e. false positive) messages.
  856.  
  857.         Does not do any processing of the AUTHENTICATE command itself."""
  858.         return self.onLogin(command_id, command, args, response)
  859.  
  860.     
  861.     def _auth_callable(self, resp):
  862.         if resp[:-1].lower() == 'user name':
  863.             return self.username
  864.         
  865.         if resp[:-1].lower() == 'password':
  866.             return self.password
  867.         
  868.  
  869.     
  870.     def add_folder_to_filter(self, folder_name):
  871.         self.filter_folders += (folder_name,)
  872.         self.save_folders_to_filter()
  873.  
  874.     
  875.     def save_folders_to_filter(self):
  876.         folders = Options.options[('imap4proxy', 'filter_folders')]
  877.         for folder in self.filter_folders:
  878.             folders += (self.current_account + '|' + folder,)
  879.         
  880.         folders = tuple(frozenset(folders))
  881.         Options.options[('imap4proxy', 'filter_folders')] = folders
  882.         Options.options.update_file(Options.optionsPathname)
  883.  
  884.     
  885.     def onSelect(self, unused, unused2, args, response):
  886.         '''Keep track of the folders that are SELECTed and EXAMINEd.'''
  887.         remainder = 0
  888.         folder_name = []
  889.         for arg in args:
  890.             folder_name.append(arg)
  891.             if arg.count('"') % 2 + remainder == 0:
  892.                 break
  893.             
  894.             remainder += arg.count('"')
  895.         
  896.         folder_name = ' '.join(folder_name)
  897.         self.add_folder_to_filter(folder_name)
  898.         return response
  899.  
  900.     
  901.     def add_welcome(self, imap_connection):
  902.         m_f_s = email.message_from_string
  903.         get_welcome = self.get_blocking_welcome_message
  904.         make_message = self.state.waitingCorpus.makeMessage
  905.         add_message = self.state.waitingCorpus.addMessage
  906.         welcomeText = get_welcome(self.username, self.serverName)
  907.         msg = m_f_s(welcomeText, _class = SEHeaderMessage)
  908.         msg_id = self.state.getNewMessageName()
  909.         msg.setId(msg_id)
  910.         corpus_msg = make_message(msg_id, msg.as_string())
  911.         add_message(corpus_msg, observer_flags = storage.NO_TRAINING_FLAG)
  912.         self.state.message_info_database.load_msg(msg)
  913.         msg.uid = -1
  914.         msg.flags = None
  915.         msg.folder_name = 'INBOX'
  916.         msg.internaldate = None
  917.         self.state.message_info_database.store_msg(msg)
  918.         self.restoreMessage(imap_connection, msg_id)
  919.  
  920.     
  921.     def onLogin(self, command_id, command, args, response):
  922.         """Spins off two separate threads that expires any old messages
  923.         in the caches, and another thread to update the message database.
  924.         Also creates an account (to hold the blocked messages) if the
  925.         server has not been seen before, and release any 'delayed'
  926.         (i.e. false positive) messages.
  927.  
  928.         Does not do any processing of the LOGIN command itself."""
  929.         SEIMAP4Proxy.onLogin(self, command_id, command, args, response, False)
  930.         if len(args) < 2:
  931.             return response
  932.         
  933.         if ' ' not in response or response.split(' ')[1].upper() != 'OK':
  934.             return response
  935.         
  936.         if command == 'LOGIN':
  937.             is_new = self._setupAccount(args[0])
  938.         elif command == 'AUTHENTICATE':
  939.             is_new = self._setupAccount(args[1])
  940.         elif not False:
  941.             raise AssertionError, 'Should never get here'
  942.         self.load_folders_to_filter()
  943.         if self.current_account in self.state.open_remote_connections:
  944.             if Options.options[('globals', 'verbose')]:
  945.                 print 'Connection to', self.current_account, 'already open, so skipping retreive + restore.'
  946.             
  947.             return response
  948.         
  949.         self.state.open_remote_connections.append(self.current_account)
  950.         i = self.imap_class(self.serverName, self.serverPort)
  951.         
  952.         try:
  953.             if command == 'LOGIN':
  954.                 self.username = args[0]
  955.                 self.password = args[1]
  956.                 connection = (self.serverName, self.serverPort, self.username, self.password, 'imap4', self.ssl)
  957.                 self.state.model_notifier.add_connection(connection)
  958.                 self.username = self.username.strip('"')
  959.                 self.password = self.password.strip('"')
  960.                 i.login(self.username, self.password)
  961.             elif command == 'AUTHENTICATE':
  962.                 print >>sys.stderr, "Can't periodic check with AUTHENTICATE"
  963.                 if args[0].lower() != 'login':
  964.                     print >>sys.stderr, "Can't handle login", args
  965.                     return None
  966.                 
  967.                 self.username = base64.decodestring(args[1])
  968.                 self.password = base64.decodestring(args[2])
  969.                 i.authenticate(args[0], self._auth_callable)
  970.         except imaplib.IMAP4.error:
  971.             e = None
  972.             print >>sys.stderr, 'Could not login to IMAP server', self.username, self.serverName
  973.             print >>sys.stderr, str(e)
  974.             
  975.             try:
  976.                 self.state.open_remote_connections.remove(self.current_account)
  977.             except ValueError:
  978.                 print >>sys.stderr, 'Connection was missing:', self.current_account
  979.  
  980.             return None
  981.  
  982.         msg = 'Logging into %s:%s\r\n' % (self.serverName, self.serverPort)
  983.         self.state.imapLogFile.write(msg)
  984.         self.state.imapLogFile.flush()
  985.         del self.password
  986.         if is_new:
  987.             self.add_welcome(i)
  988.         
  989.         if self.state.delayed_messages[self.current_account]:
  990.             restored = []
  991.             for msg_id in self.state.delayed_messages[self.current_account]:
  992.                 if Options.options[('globals', 'verbose')]:
  993.                     print >>sys.stderr, 'Restoring', msg_id
  994.                 
  995.                 if self.restoreMessage(i, msg_id):
  996.                     restored.append(msg_id)
  997.                     continue
  998.             
  999.             old = self.state.delayed_messages[self.current_account]
  1000.             for msg_id in restored:
  1001.                 del old[msg_id]
  1002.             
  1003.             self.state.delayed_messages[self.current_account] = old
  1004.         
  1005.         self.expungeMessages(i)
  1006.         update_thread = threading.Thread(target = self.updateMessageDatabase, args = (i,))
  1007.         update_thread.setDaemon(True)
  1008.         self.threads.append(update_thread)
  1009.         update_thread.start()
  1010.         return response
  1011.  
  1012.     uidnext_re = re.compile('uidnext +(\\d+)', re.IGNORECASE)
  1013.     
  1014.     def restoreMessage(self, imap_connection, msg_id):
  1015.         '''Put a message back on the server, as it was a false positive.'''
  1016.         for cache in (self.state.hamCorpus, self.state.unsureCorpus, self.state.waitingCorpus):
  1017.             msg = cache.get(msg_id)
  1018.             if msg is not None:
  1019.                 break
  1020.                 continue
  1021.         
  1022.         if msg is None:
  1023.             print >>sys.stderr, "Can't find message in cache to restore to server!"
  1024.             return False
  1025.         
  1026.         msg.load()
  1027.         msg.setId(msg_id)
  1028.         self.state.message_info_database.load_msg(msg)
  1029.         if msg.internaldate is None or msg.internaldate == [
  1030.             None]:
  1031.             date = time.localtime()
  1032.         else:
  1033.             date = 'INTERNALDATE "' + msg.internaldate + '"'
  1034.             date = imaplib.Internaldate2tuple(date)
  1035.         if not msg.flags:
  1036.             pass
  1037.         if not re.sub('\\\\Recent ?', '', '', re.IGNORECASE):
  1038.             pass
  1039.         flags = None
  1040.         
  1041.         try:
  1042.             imap_connection.select(msg.folder_name)
  1043.             (unused, response) = imap_connection.response('OK')
  1044.         except imaplib.IMAP4.error:
  1045.             e = None
  1046.             print >>sys.stderr, 'Failed to get uidnext.', msg_id
  1047.             print >>sys.stderr, str(e)
  1048.  
  1049.         for item in response:
  1050.             mo = self.uidnext_re.search(item)
  1051.             if mo:
  1052.                 next_uid = mo.group(1)
  1053.                 break
  1054.                 continue
  1055.         else:
  1056.             next_uid = None
  1057.         for flgs, dte in ((flags, date), (None, date), (flags, imaplib.Time2Internaldate(time.time())), (None, imaplib.Time2Internaldate(time.time()))):
  1058.             
  1059.             try:
  1060.                 imap_connection.append(msg.folder_name, flgs, dte, msg.as_string())
  1061.             continue
  1062.             except imaplib.IMAP4.error:
  1063.                 e = None
  1064.                 print >>sys.stderr, 'Failed putting message back on server', msg.folder_name, flgs, dte, msg_id
  1065.                 print >>sys.stderr, str(e)
  1066.                 return False
  1067.                 continue
  1068.             
  1069.  
  1070.         
  1071.         found_uid = None
  1072.         for i in xrange(20):
  1073.             if found_uid:
  1074.                 break
  1075.             
  1076.             if next_uid:
  1077.                 messageText = self._fetchAttribute(imap_connection, next_uid, 'BODY.PEEK[]', None, log = False, uid = True)
  1078.                 if messageText[0][1] == msg.as_string():
  1079.                     found_uid = next_uid
  1080.                 
  1081.             
  1082.             if not found_uid:
  1083.                 id_header = msg['Message-ID']
  1084.                 
  1085.                 try:
  1086.                     uid_func = imap_connection.uid
  1087.                     arg = '(UNDELETED HEADER Message-ID %s)' % (id_header,)
  1088.                     (unused, response) = uid_func('SEARCH', arg)
  1089.                 except imaplib.IMAP4.error:
  1090.                     e = None
  1091.                     print >>sys.stderr, 'Failed to find message.', msg_id
  1092.                     print >>sys.stderr, str(e)
  1093.  
  1094.                 matches = response[0].split()
  1095.                 if len(matches) == 1:
  1096.                     found_uid = matches[0]
  1097.                 else:
  1098.                     print >>sys.stderr, 'Found wrong number of messages', response
  1099.             
  1100.             imap_connection.noop()
  1101.         
  1102.         if not found_uid:
  1103.             print >>sys.stderr, 'Failed putting message back on server', msg.folder_name, flags, date, msg_id
  1104.             print >>sys.stderr, str(e)
  1105.             return False
  1106.         
  1107.         db_key = self._get_messageinfo_database_key(msg.folder_name, found_uid)
  1108.         msg = self._dummy_msg(db_key)
  1109.         self.state.message_info_database.load_msg(msg)
  1110.         msg.c = PERSISTENT_HAM_STRING
  1111.         self.state.message_info_database.store_msg(msg)
  1112.         self.state.imapLogFile.write('Appending %s %s %s\r\n' % (msg.folder_name, flags, date))
  1113.         self.state.imapLogFile.flush()
  1114.         return True
  1115.  
  1116.  
  1117.